from math import pi,cos,sin,atan,acos,sqrt
import matplotlib.pyplot as plt
import matplotlib.animation as animation
import numpy as np
plt.rc('font', family='Times New Roman')
plt.rc("axes", unicode_minus=False)
plt.rc("figure", figsize=(10,5))
plt.rc("savefig", dpi=400)
plt.rc("text", usetex = True)
ticklabels_style = {
    "fontname": "Times New Roman",
}

class V2D(tuple): #vector 2D
    def __init__(self,v) -> None:
        self.x = float(v[0])
        self.y = float(v[1])

    def __add__(self,other): # +
        return V2D((self.x+other.x, self.y+other.y))

    def __sub__(self,other): # -
        return V2D((self.x-other.x, self.y-other.y))
    
    def __mul__(self,other): # num *
        if isinstance(other,V2D):
            return V2D((self.x*other.x, self.y*other.y))
        elif isinstance(other,float):
            return V2D((self.x*other, self.y*other))

    def __truediv__(self,other): # /
        if isinstance(other,V2D):
            return V2D((self.x/other.x, self.y/other.y))
        elif isinstance(other,float):
            return V2D((self.x/other, self.y/other))

    def __pow__(self,other): # dot *
        return self.x*other.x+self.y*other.y

    def get_val(self): #returm tuple
        return self.x, self.y

    def dist(a,b): #return dist^2
        if b == 0:
            r = a*a
            return sqrt(r[0]+r[1])
        else:
            return V2D.dist(a-b,0)
    
    def arg(a,b):
        if b == 0:
            arg = atan(float(a.y/a.x))
            if a.x<0:
                if a.y>0:
                    arg+=pi
                elif a.y<=0:
                    arg-=pi
            return arg
        else:
            return V2D.arg(b-a,0)

a = 1
step = a/180*pi
omega = 100
dt = step/omega
length = int(360/a)

AB = 75
BC = 135
CD = 120
DA = 165
BE = 60
EF = 45
pA = V2D((0,0))
pB = V2D((0,0))
pC = V2D((0,0))
pD = V2D((DA,0))
pE = V2D((0,0))
pF = V2D((0,0))

theta = 0

B_list = []
C_list = []
E_list = []
F_list = []
for i in range(length):
    pB = pA+V2D((AB*cos(theta),AB*sin(theta)))
    BD = pB.dist(pD)
    xDB = pD.arg(pB)
    CDB = acos((BD*BD+CD*CD-BC*BC)/(2*BD*CD))
    xDC = xDB-CDB
    pC = pD+V2D((CD*cos(xDC),CD*sin(xDC)))
    xBC = pB.arg(pC)
    pE = pB+(pC-pB)*(BE/BC)
    xEF = xBC+pi/2
    pF = pE+V2D((EF*cos(xEF),EF*sin(xEF)))
    B_list.append(pB)
    C_list.append(pC)
    E_list.append(pE)
    F_list.append(pF)
    theta+=step

B_list = np.array(B_list)
C_list = np.array(C_list)
E_list = np.array(E_list)
F_list = np.array(F_list)

gs_kw = dict(width_ratios=[1.5, 1], height_ratios=[1,1,1])

fig,axd = plt.subplot_mosaic(
    [[0,'pos'],
     [0,'velo'],
     [0,'accel']], gridspec_kw=gs_kw
)
#ax = fig.add_subplot(autoscale_on=False, xlim=(-80, 200), ylim=(-80, 180))
axd      [0].set_aspect('equal')
axd      [0].set_xlim(-80, 200)
axd      [0].set_ylim(-80, 180)
axd  ['pos'].set_ylim(-0.080, 0.180)
axd  ['pos'].set_xlim(-dt*length, 0)
axd ['velo'].set_ylim(-18, 18)
axd ['velo'].set_xlim(0, 360)
axd['accel'].set_ylim(-1800, 1800)
axd['accel'].set_xlim(0, 360)
line , = axd[0].plot([], [], 'bo-', lw=1.5, ms=3)
line2, = axd[0].plot([], [], 'go-', lw=1.5, ms=3)
trace, = axd[0].plot([], [], 'co-', lw=1  , ms=0, label='trace')

axd      [0].set_xlabel('x/mm')
axd      [0].set_ylabel('y/mm')
axd  ['pos'].set_xlabel('t/s')
axd  ['pos'].set_ylabel('position/cm')
axd ['velo'].set_xlabel(r'$\varphi$/°')
axd ['velo'].set_ylabel(r'v/$cm \cdot s^{-1}$')
axd['accel'].set_xlabel(r'$\varphi$/°')
axd['accel'].set_ylabel(r'a/$cm \cdot s^{-2}$')

t_list = [(i - length)*dt for i in range(length)]
phi_list = [i*a for i in range(length)]
y_default = [0 for i in range(length)]
px_line, = axd  ['pos'].plot(t_list, y_default, 'r-', lw=1, ms=0, label='x')
py_line, = axd  ['pos'].plot(t_list, y_default, 'g-', lw=1, ms=0, label='y')
vx_line, = axd ['velo'].plot(phi_list, y_default, 'r-', lw=1, ms=0, label='x')
vy_line, = axd ['velo'].plot(phi_list, y_default, 'g-', lw=1, ms=0, label='y')
ax_line, = axd['accel'].plot(phi_list, y_default, 'r-', lw=1, ms=0, label='x')
ay_line, = axd['accel'].plot(phi_list, y_default, 'g-', lw=1, ms=0, label='y')

px_list = [0 for i in range(length)]
py_list = [0 for i in range(length)]
vx_list = [0 for i in range(length)]
vy_list = [0 for i in range(length)]
ax_list = [0 for i in range(length)]
ay_list = [0 for i in range(length)]

Fx = []
Fy = []
for i in axd:
    axd[i].legend(loc = 'upper right')

flag = 0

def animate(i):
    global flag
    thisx = [0, B_list[i,0], C_list[i,0],160,0]
    thisy = [0, B_list[i,1], C_list[i,1],0,0]
    thisx2 = [E_list[i,0], F_list[i,0]]
    thisy2 = [E_list[i,1], F_list[i,1]]
    px_list.append(F_list[i,0]/1000)
    px_list.pop(0)
    py_list.append(F_list[i,1]/1000)
    py_list.pop(0)
    vx_list.append((px_list[-1]-px_list[-2])/dt)
    vx_list.pop(0)
    vy_list.append((py_list[-1]-py_list[-2])/dt)
    vy_list.pop(0)
    ax_list.append((vx_list[-1]-vx_list[-2])/dt)
    ax_list.pop(0)
    ay_list.append((vy_list[-1]-vy_list[-2])/dt)
    ay_list.pop(0)

    Fx.append(F_list[i,0])
    Fy.append(F_list[i,1])
    if len(Fx)>=length:
        Fx.pop(0)
        Fy.pop(0)

    '''    if i == 0:
            history_x.clear()
            history_y.clear()

        history_x.appendleft(thisx[2])
        history_y.appendleft(thisy[2])
    '''
    line.set_data(thisx, thisy)
    line2.set_data(thisx2,thisy2)
    trace.set_data(Fx,Fy)
    px_line.set_ydata(px_list)
    py_line.set_ydata(py_list)
    vx_line.set_ydata(vx_list)
    vy_line.set_ydata(vy_list)
    ax_line.set_ydata(ax_list)
    ay_line.set_ydata(ay_list)
    pass
    if i == 0:
        if flag!=1:
            flag += 1
        else:
            plt.savefig('task1.png')
    #trace.set_data(history_x, history_y)
    #time_text.set_text(time_template % (i*dt))
    return line, line2, trace, px_line, py_line, vx_line, vy_line, ax_line, ay_line#, trace, time_text

def init():
    return line, line2, trace, px_line, py_line, vx_line, vy_line, ax_line, ay_line

ani = animation.FuncAnimation(
    fig, animate, range(length), interval=1, init_func=init, blit=True, repeat = True)
plt.suptitle('Linkage Dynamics Analysis - 47')
plt.tight_layout()
plt.show()
pass